import os
from numpy import array, zeros, where, log10, exp, \
                  abs, linspace, sum, cumsum, corrcoef, poly1d
from enthought.chaco.api import Plot, ArrayPlotData, \
                                ImagePlot, LinearMapper, \
                                HPlotContainer, OverlayPlotContainer, \
                                create_line_plot, PlotLabel, add_default_axes
from enthought.chaco.tools.api import PanTool, SimpleZoom, \
                                      ImageInspectorTool, ImageInspectorOverlay, \
                                      SaveTool, \
                                      RangeSelection, RangeSelectionOverlay
from enthought.chaco.tools.cursor_tool import CursorTool1D
from enthought.traits.api import Any, Bool, Int, Float, Str, Range, Button, HasTraits, \
                                DelegatesTo, Instance, Enum, Event, Trait, Tuple, \
                                Callable, List, on_trait_change
from enthought.traits.ui.api import View, Item, HGroup, VGroup, VGrid
from enthought.traits.ui.menu import OKButton, CancelButton
from enthought.enable.api import KeySpec, ColorTrait, \
                                 Component, ComponentEditor
from enthought.enable.tools.api import DragTool

# Own modules
from gprViewerUIs import wxOpenFile
from GPRformat import Picks


class StartViewerTool(HasTraits):
    file = Str
    dir = Str
    browse = Button(label='...')
    sample_min = Int(0, label="Sample Min/Max")
    sample_max = Int(-1)
    trace_min = Int(0, label="Trace Min/Max")
    trace_max = Int(-1)
    
    view = View(VGrid(Item(name = 'file'), Item(name = 'browse', show_label=False),
                Item(name = 'sample_min'), Item(name = 'sample_max', show_label=False),
                Item(name = 'trace_min'), Item(name = 'trace_max', show_label=False)),
                buttons = [OKButton, CancelButton],
                title="Choose data file and range")
    
    def __init__(self, file="", dir=os.getcwd()):
        self.file = file
        self.dir = dir
        
        
    def _file_changed(self):
        try:
            from tables import openFile, NoSuchNodeError
            h5file = openFile(os.path.join(self.dir,self.file), mode = 'r')
            try:
                self.sample_max, self.trace_max = h5file.root.data.traces.shape
            except NoSuchNodeError:
                self.sample_max, self.trace_max = h5file.root.profile.traces.shape
            h5file.close()
        except IOError:
            print "The chosen file has an unknown format."
    
    def _browse_fired(self):
        file, dir = wxOpenFile(self.dir, multi=False, filter=0)
        self.dir = dir
        self.file = file
        
class StartProcessingTool(HasTraits):
    ''' UI controlling the processing steps '''
    stack = Bool(False)
    stack_no = Int(1)
    remove_gating = Bool(True)
    linear_gain = Bool(False)
    lin_gain_val = Float(20.0)
    compensate_spreading = Bool(True)
    comp_spread_val = Float(2.0)
    normalize = Bool(True)
    decibel = Bool(True)
    correct_topography = Bool(False)
    
    spacer = Str() # keeps some space where other items have an input field
    
    check_group = VGrid(
                        Item(name="stack"), Item(name="stack_no", show_label=False, enabled_when="stack==True"),
                        Item(name="remove_gating"), Item(name="spacer", show_label=False, style="readonly"),
                        Item(name="linear_gain"), Item(name="lin_gain_val", show_label=False, enabled_when="linear_gain==True"),
                        Item(name="compensate_spreading"), Item(name="comp_spread_val", show_label=False, enabled_when="compensate_spreading==True"),
                        Item(name="normalize"), Item(name="spacer", show_label=False, style="readonly"),
                        Item(name="decibel"), Item(name="spacer", show_label=False, style="readonly"),
                        Item(name="correct_topography"), Item(name="spacer", show_label=False, style="readonly"),
                        )
    
    view = View(
                check_group,
                buttons = [OKButton, CancelButton],
                title = "Choose processing steps..."
                )
    
class StartConvertTool(HasTraits):
    filein = Str
    fbrowse = Button(label='...')
    dirout = Str
    dbrowse = Button(label='...')
    timens = Float(2000.0)
    stacks = Int(1)
    combine = Bool(False)
    filerange = Range(0, 9999)   
    
    view = View(HGroup(Item(name = 'filein', label='Input file'), Item(name = 'fbrowse', show_label=False)),
                HGroup(Item(name = 'dirour', label='Output dir'), Item(name = 'dbrowse', show_label=False)),
                Item(name = 'timens'), Item(name = 'stacks'),
                Item(name = 'combine'), Item(name = 'filerange'),
                buttons = [OKButton, CancelButton],
                title="Convert HUBRA file")
    
    def _fbrowse_fired(self):
        self.file, self.dir = wxOpenFile(multi=False)
        
    def _dbrowse_fired(self):
        pass
#        self.file, self.dir = wxOpenFile(multi=False) 


class ShiftSaveTool(SaveTool):
    """ This tool allows the user to press Ctrl+Shift+S to save a snapshot image of
    the plot component.
    """
    def normal_key_pressed(self, event):
        """ Handles a key-press when the tool is in the 'normal' state.
        
        Saves an image of the plot if the keys pressed are Control and S.
        """
        if self.component is None:
            return

        if event.character == "s" and event.control_down and event.shift_down:
            if os.path.splitext(self.filename)[-1] == ".pdf":
                self._save_pdf()
            else:
                self._save_raster()
            event.handled = True
        return
    

class PickTool(ImageInspectorTool):
    """ A tool that captures the color and underlying values of an image plot.
    """

    # This event fires whenever the mouse moves over a new image point.
    # Its value is a dict with a key "color_value", and possibly a key
    # "data_value" if the plot is a color-mapped image plot.
    new_value = Event

    # Indicates whether overlays listening to this tool should be visible.
    visible = Bool(False)
    
    pointer = "bullseye"
    
    # Stores the last mouse position.  This can be used by overlays to
    # position themselves around the mouse.
    last_mouse_position = Tuple

    # This key will show and hide any ImageInspectorOverlays associated
    # with this tool.
    enable_key = KeySpec('p')
    save_key = KeySpec('s')
    reset_key = KeySpec('r')

   
    # Stores the value of self.visible when the mouse leaves the tool,
    # so that it can be restored when the mouse enters again.
#    _old_visible = Trait(None, Bool(True))
    _old_visible = Enum(None, True, False)
    _enabled = Bool(False)
    # Container for the picked data
    # Imports from GPRformat.py
    picks = Picks()
    triggered = Bool(False)
    data = None
    
    def enable(self):
        self.visible = not self.visible
#        if self.visible:
#            event.window.set_pointer(self.pointer)
#        else:
#            event.window.set_pointer("arrow")
    
    def normal_key_pressed(self, event):    
        if self.enable_key.match(event):
            self.visible = not self.visible
            if self.visible:
                event.window.set_pointer(self.pointer)
            else:
                event.window.set_pointer("arrow")
        elif self.visible and self.save_key.match(event):
#            try:
            self.picks.save('test_pck.csv')
            print "Picks save to file"
#            except:
        elif self.visible and self.reset_key.match(event):
            self.reset_picks()


    def normal_left_down(self, event):
        """ Handles the left mouse button being clicked. """
        if self.visible:
            plot = self.component
            if plot is not None:
                if isinstance(plot, ImagePlot):
                    ndx = plot.map_index((event.x, event.y))
                    x_index, y_index = ndx
                    self.picks.append(self.data.tracendx[x_index],
                                      self.data.t[y_index],
                                      self.data.coords.x[x_index],
                                      self.data.coords.y[x_index],
                                      self.data.coords.elev[x_index])
                    # needs to change at each click so the on_trait_change callback in gprViewer.py is fired
                    self.triggered = not self.triggered
                    
            return
        else:
            pass
    
    def normal_mouse_move(self, event):
        """ Handles the mouse being moved.
        
        Fires the **new_value** event with the data (if any) from the event's 
        position.
        """
        plot = self.component
        if plot is not None:
            if isinstance(plot, ImagePlot):
                ndx = plot.map_index((event.x, event.y))
                if ndx == (None, None):
                    self.new_value = None
                    return

                x_index, y_index = ndx
                image_data = plot.value
                if hasattr(plot, "_cached_mapped_image") and plot._cached_mapped_image is not None:
                    try:
                        self.new_value = dict(indices=ndx,
                                              map_values=[int(self.data.tracendx[x_index]), self.data.t[y_index]],
                                              data_value=image_data.data[y_index, x_index],
                                              color_value=plot._cached_mapped_image[y_index, x_index])
                    except IndexError:
                        try:
                            self.new_value = dict(indices=ndx,
                                                  map_values=[int(self.data.tracendx[x_index]), self.data.r[y_index]],
                                                  data_value=image_data.data[y_index, x_index],
                                                  color_value=plot._cached_mapped_image[y_index, x_index])
                        except IndexError:   
                            try:
                                self.new_value = dict(indices=ndx,
                                                      map_values=[int(self.data.tracendx[x_index]), self.data.e[y_index]],
                                                      data_value=image_data.data[y_index, x_index],
                                                      color_value=plot._cached_mapped_image[y_index, x_index])
                            except IndexError:
                                print "Unknown index format!"
                                pass    
                    
                        
                else:
                    self.new_value = \
                        dict(indices=ndx,
                             color_value=image_data.data[y_index, x_index])
                self.last_mouse_position = (event.x, event.y)
        return
    
    def reset_picks(self):
        self.picks = Picks()
        self.component.request_redraw()
        print "Pick container reset"
        
    def get_data(self, data):
        # Maybe implement to get only time, trace and coordinate from data, not the signal
        # since it is very large.
        pass
        
    

class PickToolOverlay(ImageInspectorOverlay):
    """ An overlay that displays a box containing values from an 
    ImageInspectorTool instance.
    """
    # An instance of ImageInspectorTool; this overlay listens to the tool
    # for changes, and updates its displayed text accordingly.
    image_inspector = Any

    # Anchor the text to the mouse?  (If False, then the text is in one of the
    # corners.)  Use the **align** trait to determine which corner.
    tooltip_mode = Bool(True)

    # The default state of the overlay is invisible (overrides PlotComponent).
    visible = False

    # Whether the overlay should auto-hide and auto-show based on the
    # tool's location, or whether it should be forced to be hidden or visible.
    visibility = Enum("auto", True, False)

    def _image_inspector_changed(self, old, new):
        if old:
            old.on_trait_event(self._new_value_updated, 'new_value', remove=True)
            old.on_trait_change(self._tool_visible_changed, "visible", remove=True)
        if new:
            new.on_trait_event(self._new_value_updated, 'new_value')
            new.on_trait_change(self._tool_visible_changed, "visible") 
            self._tool_visible_changed()
        
    def _new_value_updated(self, event):
        if event is None:
            self.text = ""
            if self.visibility == "auto":
                self.visible = False
            return
        elif self.visibility == "auto":
            self.visible = True

        if self.tooltip_mode:
            self.alternate_position = self.image_inspector.last_mouse_position
        else:
            self.alternate_position = None
        
        d = event
        newstring = ""
        if 'indices' in d:
            newstring += '(%d, %d)' % d['indices'] + '\n'
        if 'map_values' in d:
            newstring += "(%i, %5.2f)" % (d['map_values'][0], d['map_values'][1]) + "\n"
        if 'data_value' in d:
            newstring += str(d['data_value'])

        self.text = newstring
        self.component.request_redraw()

# ADD REQUIRED MODULES
class PointDraggingTool(DragTool):
    
    component = Instance(Component)

    # The pixel distance from a point that the cursor is still considered
    # to be 'on' the point
    threshold = Int(5)

    # The index of the point being dragged
    _drag_index = Int(-1)

    # The original dataspace values of the index and value datasources
    # corresponding to _drag_index
    _orig_value = Tuple

    def is_draggable(self, x, y):
        # Check to see if (x,y) are over one of the points in self.component
        if self._lookup_point(x, y) is not None:
            return True
        else:
            return False

    def normal_mouse_move(self, event):
        plot = self.component
        
        ndx = plot.map_index((event.x, event.y), self.threshold)
        if ndx is None:
            if plot.index.metadata.has_key('selections'):
                del plot.index.metadata['selections']
        else:
            plot.index.metadata['selections'] = [ndx]

        plot.invalidate_draw()
        plot.request_redraw()


    def drag_start(self, event):
        plot = self.component
        ndx = plot.map_index((event.x, event.y), self.threshold)
        if ndx is None:
            return
        self._drag_index = ndx
        self._orig_value = (plot.index.get_data()[ndx], plot.value.get_data()[ndx])

    def dragging(self, event):
        plot = self.component

        data_x, data_y = plot.map_data((event.x, event.y))

        plot.index._data[self._drag_index] = data_x
        plot.value._data[self._drag_index] = data_y
        plot.index.data_changed = True
        plot.value.data_changed = True
        plot.request_redraw()

    def drag_cancel(self, event):
        plot = self.component
        plot.index._data[self._drag_index] = self._orig_value[0]
        plot.value._data[self._drag_index] = self._orig_value[1]
        plot.index.data_changed = True
        plot.value.data_changed = True
        plot.request_redraw()

    def drag_end(self, event):
        plot = self.component
        if plot.index.metadata.has_key('selections'):
            del plot.index.metadata['selections']
        plot.invalidate_draw()
        plot.request_redraw()

    def _lookup_point(self, x, y):
        """ Finds the point closest to a screen point if it is within self.threshold
        
        Parameters
        ==========
        x : float
            screen x-coordinate
        y : float
            screen y-coordinate

        Returns
        =======
        (screen_x, screen_y, distance) of datapoint nearest to the input *(x,y)*.
        If no data points are within *self.threshold* of *(x,y)*, returns None.
        """

        if hasattr(self.component, 'get_closest_point'):
            # This is on BaseXYPlots
            return self.component.get_closest_point((x,y), threshold=self.threshold)

        return None



class TraceSelectTool(SimpleZoom):
    """ Selects a range along the trace axis.

    The user left-click-drags to select a region.
    Certain keyboard keys are mapped to performing actions as well.
    """

    # The selection mode:
    #
    # range:
    #   Select a range across a single index or value axis.
    # box:
    #   Perform a "box" selection on two axes.
    tool_mode = Enum("range", "box")

    # Is the tool always "on"? If True, left-clicking always initiates
    # a zoom operation; if False, the user must press a key to enter zoom mode.
    always_on = Bool(False)

    # Defines a meta-key, that works with always_on to set the zoom mode. This
    # is useful when the zoom tool is used in conjunction with the pan tool.
    always_on_modifier = Enum(None, 'shift', 'control', 'alt')

    #-------------------------------------------------------------------------
    # Zoom control
    #-------------------------------------------------------------------------

    # The axis to which the selection made by this tool is perpendicular. This
    # only applies in 'range' mode.
    axis = Enum("index", "value")

    #-------------------------------------------------------------------------
    # Interaction control
    #-------------------------------------------------------------------------

    # Enable the mousewheel for zooming?
    enable_wheel = Bool(False)

    # The mouse button that initiates the drag.  If "None", then the tool
    # will not respond to drag.  (It can still respond to mousewheel events.)
    drag_button = Enum("left", "right", None)

    # Conversion ratio from wheel steps to zoom factors.
    wheel_zoom_step = Float(1.0)

    # The key press to enter zoom mode, if **always_on** is False.  Has no effect
    # if **always_on** is True.
    enter_zoom_key = Instance(KeySpec, args=("t",))

    # The key press to leave zoom mode, if **always_on** is False.  Has no effect
    # if **always_on** is True.
    exit_zoom_key = Instance(KeySpec, args=("t",))

    # Disable the tool after the zoom is completed?
    disable_on_complete = Bool(True)

    # The minimum amount of screen space the user must select in order for
    # the tool to actually take effect.
    minimum_screen_delta = Int(1)

    #-------------------------------------------------------------------------
    # Appearance properties (for Box mode)
    #-------------------------------------------------------------------------

    # The pointer to use when drawing a zoom box.
    pointer = "size right"

    # The color of the selection box.
    color = ColorTrait("red")

    # The alpha value to apply to **color** when filling in the selection
    # region.  Because it is almost certainly useless to have an opaque zoom
    # rectangle, but it's also extremely useful to be able to use the normal
    # named colors from Enable, this attribute allows the specification of a
    # separate alpha value that replaces the alpha value of **color** at draw
    # time.
    alpha = Trait(0.4, None, Float)

    # The color of the outside selection rectangle.
    border_color = ColorTrait("red")

    # The thickness of selection rectangle border.
    border_size = Int(1)

    # The possible event states of this zoom tool.
    event_state = Enum("normal", "selecting")
    
    # List that contains the traces selected
    remove_list = List
    
    # On trait change confirmed the selected traces will be removed
    confirmed = Bool(False)
    
    #------------------------------------------------------------------------
    # Key mappings
    #------------------------------------------------------------------------

    # The key that cancels the zoom and resets the view to the original defaults.
    cancel_zoom_key = Instance(KeySpec, args=("Esc",))
    
    # The key that confirms the deletion.
    confirm_key = Instance(KeySpec, args=("a",))
    

    #------------------------------------------------------------------------
    # Private traits
    #------------------------------------------------------------------------

    # If **always_on** is False, this attribute indicates whether the tool
    # is currently enabled.
    _enabled = Bool(False)

    # the original numerical screen ranges
    _orig_low_setting = Trait(None, Tuple, Float, Str)
    _orig_high_setting = Trait(None, Tuple, Float, Str)

    # The (x,y) screen point where the mouse went down.
    _screen_start = Trait(None, None, Tuple)

    # The (x,,y) screen point of the last seen mouse move event.
    _screen_end = Trait(None, None, Tuple)

    def __init__(self, component=None, *args, **kw):
        # Support AbstractController-style constructors so that this can be
        # handed in the component it will be overlaying in the constructor
        # without using kwargs.
        self.component = component
        super(SimpleZoom, self).__init__(*args, **kw)
        self._reset_state_to_current()
        if self.tool_mode == "range":
            mapper = self._get_mapper()
            self._orig_low_setting = mapper.range.low_setting
            self._orig_high_setting = mapper.range.high_setting
        else:
            x_range = self.component.x_mapper.range
            y_range = self.component.y_mapper.range
            self._orig_low_setting = (x_range.low_setting, y_range.low_setting)
            self._orig_high_setting = \
                (x_range.high_setting, y_range.high_setting)
        component.on_trait_change(self._reset_state_to_current,
                                  "index_data_changed")
        return
    
    
    def normal_key_pressed(self, event):
        """ Handles a key being pressed when the tool is in 'normal' state.

        If the tool is not always on, this method handles turning it on and
        off when the appropriate keys are pressed. Also handles keys to
        manipulate the tool history.
        """
        if not self.always_on:
            if not self._enabled and self.enter_zoom_key.match(event):
                if self.component.active_tool in (None, self):
                    self.component.active_tool = self
                    self._enabled = True
                    event.window.set_pointer(self.pointer)
                else:
                    self._enabled = False
                return
            elif self._enabled and self.exit_zoom_key.match(event):
                self._enabled = False
                event.window.set_pointer("arrow")
                return
            elif self.confirm_key.match(event):
                self.confirmed = not self.confirmed
                return

        self._history_handle_key(event)

        if event.handled:
            self.component.request_redraw()
        return
        
    
    
    def _start_select(self, event):
        """ Starts selecting the zoom region
        """
        if self.component.active_tool in (None, self):
            self.component.active_tool = self
        else:
            self._enabled = False
        self._screen_start = (event.x, event.y)
        self._screen_end = None
        self.event_state = "selecting"
        event.window.set_pointer(self.pointer)
        event.window.set_mouse_owner(self, event.net_transform())
        self.selecting_mouse_move(event)
        return
    
    
    def _end_select(self, event):
        """ Ends selection of the zoom region, adds the new zoom range to
        the zoom stack, and does the zoom.
        """
        self._screen_end = (event.x, event.y)

        start = array(self._screen_start)
        end = array(self._screen_end)

        if sum(abs(end - start)) < self.minimum_screen_delta:
            self._end_selecting(event)
            event.handled = True
            return

        if self.tool_mode == "range":
            axis = self._determine_axis()
            low = self.component.map_index(self._screen_start)[axis]
            high = self.component.map_index(self._screen_end)[axis]
            if low > high:
                low, high = high, low
        else:
            low, high = self._map_coordinate_box(self._screen_start, self._screen_end)

        new_trace_range = (low, high)
#        self._append_state(new_trace_range)
        self.remove_list.append(new_trace_range)
        print "Press a to remove trace-range:", new_trace_range
        self._end_selecting(event)
        event.handled = True
        return
    
        
    
    

class TrackPlot(HasTraits):
    from enthought.chaco.default_colormaps import jet#*
    container = Instance(OverlayPlotContainer)
    traits_view = View(
                       Item('container',editor=ComponentEditor(), show_label=False),
                       width=500, height=500, resizable=True, title="Track Plot"
                       )
    
    cmap = Trait(jet, Callable)
    
    def __init__(self, x, y, elev):
        super(TrackPlot, self).__init__()
        self.plotdata = ArrayPlotData(x=x, y=y, elev=elev,
                                 pointx=array([]), pointy=array([]))
        plot = Plot(self.plotdata,
                    padding_left=80, padding_bottom=60,
                    padding_right=20, padding_top=20,)
        scatterplot = plot.plot(("x", "y", "elev"), type="cmap_scatter", color_mapper=self.cmap(None),
                                name="track_plot")[0]
        add_default_axes(plot, orientation = 'normal',
                         vtitle = 'Y coordinate [m]', htitle = 'X coordinate [m]')
        scatterplot.fill_alpha = 0.75
        scatterplot.outline_color = "none"
        
        pos_circle = plot.plot(("pointx", "pointy"), type="scatter", name="point_plot")[0]
        pos_circle.marker = "circle"
        pos_circle.marker_size = 20
        pos_circle.color = (0.0, 0.0, 0.4, 0.5)
        pos_cross = plot.plot(("pointx", "pointy"), type="scatter", name="point_plot")[0]
        pos_cross.marker = "dot"
        pos_cross.marker_size = 2
        pos_cross.line_width = 3.0
        pos_cross.color = (0.0, 0.0, 0.4, 1.0)
        self.plotdata.set_data("pointx", array([x[2]]))
        self.plotdata.set_data("pointy", array([y[2]]))
        
        
        plot.tools.append(PanTool(plot))
        zoom = SimpleZoom(component=plot, tool_mode="box", drag_button="right", always_on=True)
        plot.overlays.append(zoom)
        
        container = OverlayPlotContainer()
        container.add(plot)
        self.container = container

class ReflectivityPlot(HasTraits):
    plot = Instance(Plot)
    traits_view = View(
        Item('plot',editor=ComponentEditor(), show_label=False),
        width=500, height=500, resizable=True, title="Reflectivity vs Depth Plot")
    
    def __init__(self, x, y):
        super(ReflectivityPlot, self).__init__()
        plotdata = ArrayPlotData(x=x, y=y)
        plot = Plot(plotdata)
        plot.plot(("x", "y"), type="scatter")

        plot.tools.append(PanTool(plot))
        zoom = SimpleZoom(component=plot, tool_mode="box", drag_button="right", always_on=True)
        plot.overlays.append(zoom)
        
        self.plot = plot


class SelectCursorTool(CursorTool1D):
    " Simply adds a key_press event to CursorTool1D "
    select_key = KeySpec('x')
    selected = Bool(False)
    
    def normal_key_pressed(self, event):    
        if self.select_key.match(event):
            self.selected = not self.selected


class DecayPlot(HasTraits):
    hcontainer = Instance(HPlotContainer)
    extinction_fit = Str("")
    penetration_depth = Float()
    root1 = Float()
    root2 = Float()
    cmed = Float(0.23, label="Medium velocity [m/ns]")
    
    traits_view = View(
                       VGroup(
                              Item('hcontainer', editor=ComponentEditor(), show_label=False),
                              VGroup(
                                     HGroup(
                                            Item('penetration_depth', style='readonly', format_str="%.2g"),
                                            Item('extinction_fit', style='readonly'),
                                            Item('cmed', style='readonly', format_str="%.3g")
                                            ),
#                                     Item('ref_power', label="Reference power"),
                                     show_border=True),
                                     HGroup(
                                            Item('root1', style='readonly', format_str="%.2g"),
                                            Item('root2', style='readonly', format_str="%.2g"),
                                            ),
                              ),
                              width=500, height=800, resizable=True, title="Extinction Fit")
    

    def __init__(self, x, y, cmed, title=""):
        super(DecayPlot, self).__init__()
        self.cmed = cmed
        self.x = x*self.cmed/2
        y = 10**(y/10)
        total_power = sum(y)
        z_power = cumsum(y)
        self.y = (total_power-z_power)/total_power
#        print total_power
#        from pylab import plot, show
#        plot(z_power)
#        show()
        lineplot = create_line_plot((self.x, self.y), index_bounds = None, value_bounds = None,
                         orientation = 'v', color = 'blue', width = 2.0, dash = 'solid',
                         value_mapper_class = LinearMapper,
                         bgcolor = 'transparent', border_visible = True,
                         add_grid = False, add_axis = False, index_sort = 'ascending')
        add_default_axes(lineplot, orientation = 'flipped',
                         vtitle = 'Depth [m]', htitle = 'Two-way Extinction')
        lineplot.origin = "top left"
        lineplot.tools.append(PanTool(lineplot, drag_button="right"))
        zoom = SimpleZoom(component=lineplot, tool_mode="box", drag_button="right", always_on=False)
        lineplot.overlays.append(zoom)
        
        self.scatterdata = ArrayPlotData(locx1=array([]), locy1=array([]),
                                         locx2=array([]), locy2=array([]),
                                         decayx=array([]), decayy=array([]))

        self.scatterplot = Plot(self.scatterdata, default_origin="top left",
                                padding=lineplot.padding,
                                orientation=lineplot.orientation)
        self.scatterplot.plot(("locx1", "locy1"), type="scatter", color="blue", 
                         marker="cross", marker_size=10, line_width=1.5)
        self.scatterplot.plot(("locx2", "locy2"), type="scatter", color="black", 
                         marker="cross", marker_size=10, line_width=1.5)
        self.scatterplot.plot(("decayx", "decayy"), type="line",
                              color="black", line_width=2.0, line_style="long dash")
        
        self.scatterplot.index_range = lineplot.index_range
        self.scatterplot.value_range = lineplot.value_range
        self.label = PlotLabel("",
                               component=self.scatterplot,
                               font = "swiss 12", #border_color='red', border_visble=True, bgcolor="red",
                               overlay_position="inside bottom",
                               padding=20)
        self.scatterplot.overlays.append(self.label)
        self.scatterplot.x_axis.visible = False
        self.scatterplot.y_axis.visible = False
        
        lineplot.padding_right = 0
        self.scatterplot.padding = lineplot.padding
        
        ocontainer = OverlayPlotContainer(padding_left=60, padding_bottom=60, padding_top=60, padding_right=20)
        ocontainer.add(lineplot)
        ocontainer.add(self.scatterplot)
        ocontainer.overlays.append(PlotLabel(title,
                              component=ocontainer,
                              font = "swiss 16",
                              overlay_position="top"))
        ocontainer.tools.append(SaveTool(ocontainer, always_on=True, filename="power_ratio.png"))
        self.hcontainer = HPlotContainer(width=500, height=800)
        self.hcontainer.add(ocontainer)
        self.decay_fit()
        
        
    def decay_fit(self):
        from scipy.optimize import leastsq
        fitfunc = lambda p: exp(-2*p[0]*self.x-p[1]*self.x**2) # Target function
#        fitfunc = lambda p: exp(-(p[0]+p[1]*self.x)) # Target function
#        fitfunc = lambda p: exp(-(2*p*self.x)) # Target function
        errfunc = lambda p: fitfunc(p) - self.y # Distance to the target function
        p = [0.2, 0.2] # Initial guess for the parameter
#        p = 0.02
        param, success = leastsq(errfunc, p)
        print param, "-", success
        self.scatterdata.set_data("decayx", self.x)
        self.scatterdata.set_data("decayy", fitfunc(param)) 
        self.extinction_fit = "%.3e + %.3e z" % (param[0], param[1])
        pn = poly1d([param[1]*0.5,
                         param[0],
                         -1])
        [self.root1, self.root2] = pn.r
        if self.root1>0.0 and self.root2>0:
            if self.root1>self.root2:
                self.penetration_depth = self.root2
            else:
                self.penetration_depth = self.root1
        else:
            if self.root1>self.root2:
                self.penetration_depth = self.root1
            else:
                self.penetration_depth = self.root2
        try:
            ndx1 = where(self.y<=exp(-2))[0][0]
            ndx2 = where(fitfunc(param)<=exp(-2))[0][0]
#            self.penetration_depth = self.x[ndx1]
            self.scatterdata.set_data("locy1", array([self.y[ndx1]]))
            self.scatterdata.set_data("locx1", array([self.x[ndx1]]))
            self.scatterdata.set_data("locy2", array([fitfunc(param)[ndx2]]))
            self.scatterdata.set_data("locx2", array([self.x[ndx2]]))
        except IndexError:
            print "WARNING: The curve does not decay by 1/e in the given time window!"
        cc = corrcoef(self.y, fitfunc(param))[0,1]
        self.label.text = "R = %2.2f" % cc
#        from pylab import plot, show
#        plot(self.x, fitfunc(param)-self.y)
#        show()
        
         

class TracePlot(HasTraits):
    hcontainer = Instance(HPlotContainer)
#    cursor = Instance(SelectCursorTool)
#    cursor_pos = DelegatesTo('cursor', prefix='current_position')
    top_reflection = Int(10)
    noise_floor = Int(30)
    path_length = Int(40)
    total_power = Float()
    ref_power = Float()
    skin_depth = Float()
    data_range = Tuple()
    cmed = Float(0.24, label="Medium velocity [m/ns]")
#    top_reflection = Range(0,100,10)
#    noise_floor = Range(0,100,30)
    traits_view = View(
                       VGroup(
                              Item('hcontainer', editor=ComponentEditor(), show_label=False),
                              VGroup(
#                                     Item('cursor_pos', width=300, style="readonly",
#                                          label="Cursor position:"),
                                     HGroup(
                                            Item('skin_depth', style='readonly', format_str="%5.3f"),
                                            Item('total_power', style='readonly', format_str="%5.3f"),
                                            Item('cmed', style='readonly', format_str="%5.3f")),
#                                     Item('ref_power', label="Reference power"),
                                     show_border=True),
                              Item('top_reflection'),
                              Item('noise_floor'),
                              Item('path_length')
                              ),
                              width=500, height=800, resizable=True, title="Attenuation")
    selection_list = []
    regtext = ""
    decaytext = ""
    

    def __init__(self, x, y, title=""):
        super(TracePlot, self).__init__()
        self.x = x*self.cmed/2
        self.y = y
#        self.top_reflection._high = len(x)
#        self.noise_floor._high = len(x)
        lineplot = create_line_plot((self.x, self.y), index_bounds = None, value_bounds = None,
                         orientation = 'v', color = 'blue', width = 1.0, dash = 'solid',
                         value_mapper_class = LinearMapper,
                         bgcolor = 'transparent', border_visible = True,
                         add_grid = False, add_axis = True, index_sort = 'ascending')
        lineplot.origin = "top left"
        lineplot.tools.append(PanTool(lineplot, drag_button="right"))
        zoom = SimpleZoom(component=lineplot, tool_mode="box", drag_button="right", always_on=False)
        lineplot.overlays.append(zoom)
        self.range_selection = RangeSelection(lineplot, left_button_selects = True)
        lineplot.active_tool = self.range_selection
        lineplot.overlays.append(RangeSelectionOverlay(component=lineplot))
        self.lineplot = lineplot
#        self.cursor = SelectCursorTool(component=lineplot, drag_button="left", color='lightblue')
#
#        lineplot.overlays.append(self.cursor)
#        self.cursor.on_trait_change(self.selected, 'selected')
        
        self.scatterdata = ArrayPlotData(power=array([]), time=array([]),
                                         regx=array([]), regy=array([]),
                                         decayx=array([]), decayy=array([]))

        self.scatterplot = Plot(self.scatterdata, default_origin="top left",
                                padding=lineplot.padding,
                                orientation=lineplot.orientation)
        self.scatterplot.plot(("power", "time"), type="scatter", color="red", 
                         marker="dot", marker_size=2)[0]
        self.scatterplot.plot(("regx", "regy"), type="line", color="red")
        self.scatterplot.plot(("decayx", "decayy"), type="line", color="black")
        
        self.scatterplot.index_range = lineplot.index_range
        self.scatterplot.value_range = lineplot.value_range
        self.label = PlotLabel("",
                               component=self.scatterplot,
                               font = "swiss 12",
                               overlay_position="inside bottom",
                               padding=20)
        self.scatterplot.overlays.append(self.label)
        self.scatterplot.x_axis.visible = False
        self.scatterplot.y_axis.visible = False
        
        lineplot.padding_right = 0
        self.scatterplot.padding = lineplot.padding
        
        ocontainer = OverlayPlotContainer(padding=40)
        ocontainer.add(lineplot)
        ocontainer.add(self.scatterplot)
        ocontainer.overlays.append(PlotLabel(title,
                              component=ocontainer,
                              font = "swiss 16",
                              overlay_position="top"))
        ocontainer.tools.append(SaveTool(ocontainer, always_on=True, filename="lin_fit.png"))
        self.hcontainer = HPlotContainer(width=500, height=800)
        self.hcontainer.add(ocontainer)
        self.top_reflection = 142
        self.noise_floor = 4500
        
    @on_trait_change("range_selection.selection")
    def data_range_ch(self):
        print self.lineplot.metadata["selections"]
        print self.range_selection.selection
    
    def ke_constant(self, x, y):
        from scipy.optimize import leastsq
#        loge1 = 10*log10(exp(1))
        loge2 = 20*log10(exp(1))
        # Fitting function
        fitfunc = lambda p: p[0] -loge2*p[1]*x
        # Distance to the target function
        errfunc = lambda p: fitfunc(p) - y
        # Initial guess for the parameter
        p = [-32.0, 7.0e-01]
        
        fit_params, success = leastsq(errfunc, p)
#        print fit_params, "-", success
        
        if fit_params is not None:
            # draw the label overlay                 
            self.regtext = "P0=%.5f, ke=%.5f " % (fit_params[0], fit_params[1])
            self.label.text = self.regtext + "\n" + self.decaytext    
            pn = poly1d([fit_params[1],
                         -1])  
            self.skin_depth = pn.r[0]
            regx = linspace(self.scatterplot.index_range.low, self.scatterplot.index_range.high, len(x))
            regy = fit_params[0] -loge2*fit_params[1]*regx
            self.scatterdata.set_data("regx", regx)
            self.scatterdata.set_data("regy", regy) 
    
    
    def ke_linear(self, x, y):
        from scipy.optimize import leastsq
        loge1 = 10*log10(exp(1))
        loge2 = 20*log10(exp(1))
        # Fitting function
        fitfunc = lambda p: p[0] - loge2*p[1]*x-loge1*p[2]*x**2
        # Distance to the target function
        errfunc = lambda p: fitfunc(p) - y
        # Initial guess for the parameter
        p = [-32.0, 7.0e-03, -2.0e-03]
        
        fit_params, success = leastsq(errfunc, p)
        print fit_params, "-", success

        if fit_params is not None:
            # draw the label overlay 
            self.regtext = "RED: k0=%.5f, A=%.5f, P0=%.5f" % (fit_params[1], fit_params[2], fit_params[0])
            self.label.text = self.regtext + "\n" + self.decaytext    
                       
            pn = poly1d([fit_params[2]*0.5,
                         fit_params[1],
                         -1])
            [a1, a2] = pn.r
            if a1>0.0 and a2>0:
                if a1>a2:
                    self.skin_depth = a2
                else:
                    self.skin_depth = a1
            else:
                if a1>a2:
                    self.skin_depth = a1
                else:
                    self.skin_depth = a2
            print "skin depth +/-", pn.r
            regx = linspace(self.scatterplot.index_range.low, self.scatterplot.index_range.high, len(x))
            regy = fit_params[0] - loge2*fit_params[1]*regx-loge1*fit_params[2]*regx**2
            self.scatterdata.set_data("regx", regx)
            self.scatterdata.set_data("regy", regy)
    
    
    @on_trait_change("top_reflection, noise_floor, path_length")
    def decay_upper(self):
        print self.range_selection.selection, "Need to subclass RangeSelection.set_selection"
        if self.noise_floor>self.top_reflection:
            pathlength = self.path_length
            pathsamples = int((self.noise_floor-self.top_reflection) / pathlength)
            peaksx = zeros((pathsamples,))
            peaksy = zeros((pathsamples,)) 
            for n in range(pathsamples):
                top = self.top_reflection+n*pathlength
                bot = self.top_reflection+((n+1)*pathlength)
                pathy = self.y[top:bot]
                pathx = self.x[top:bot]
                peaksy[n] = pathy.max()
                peaksx[n] = pathx[where(self.y[top:bot]==self.y[top:bot].max())[0]]
            if len(peaksx) < 2 or len(peaksy) < 2:
                fit_params = None
            else:
#                self.ke_constant(peaksx, peaksy) 
                self.ke_linear(peaksx, peaksy)       
    
            self.scatterdata.set_data("power", peaksx)
            self.scatterdata.set_data("time", peaksy)
        else:
            pass

            
#    @on_trait_change("top_reflection, noise_floor")        
    def decay_plot(self):
        dp = DecayPlot(self.x[self.top_reflection:self.noise_floor],
                                   self.y[self.top_reflection:self.noise_floor],
                                   "P(zmax) - P(z)")
        dp.edit_traits()
    
    
#    def selected(self):
#        self.selection_list.append(self.cursor.current_position)
#        y = array(self.selection_list)[:,1]
#        x = array(self.selection_list)[:,0]
#        if len(x) < 2 or len(y) < 2:
#            fit_params = None
#            centroid = None
#        else:
#            fit = polyfit(x,y,1)
#            fit_params = tuple(fit)
#            centroid = (sum(x)/len(x)), (sum(y)/len(y))
#        if fit_params is not None:
#            # draw the label overlay 
#            if fit_params[1] < 0:
#                operator = "-"
#            else:
#                operator = "+"
#                
#            self.regtext = "RED: %.5fx "%fit_params[0] + operator + \
#                               " %.2f" % abs(fit_params[1])
#            self.label.text = self.regtext + "\n" + self.decaytext    
#              
#            regx = linspace(self.scatterplot.index_range.low, self.scatterplot.index_range.high, len(x))
#            regy = polyval(fit, regx)
#            self.scatterdata.set_data("regx", regx)
#            self.scatterdata.set_data("regy", regy)        
#
#        self.scatterdata.set_data("power", x)
#        self.scatterdata.set_data("time", y)     

#    def _ref_power_changed(self):
##        needs to address the data of the line plot not the scatter plot
#        self.y = self.y-self.ref_power
#        self.scatterdata.set_data("power", self.y)



class SkinDepthTool(HasTraits):
    traces = Int(0)
    refndxmin = Int(130)
    refndxmax = Int(225)
    
    def _init_(self, data, traces, time):
        self.data = data
        self.traces = traces
        self.time = time
        self.refndx = zeros((self.traces,1))
        self.max = zeros((self.traces,1))
        self.skindx = zeros((self.traces,1))
        self.skindepth = zeros((self.traces,1))
        self.refdepth = zeros((self.traces,1))
        
    def get_skin_depth(self):
        for n in range(self.traces):
            trace = self.data[:,n]
            self.max[n,0] = trace[self.refndxmin:self.refndxmax].max()
            self.refndx[n,0] = where(trace[self.refndxmin:self.refndxmax]==trace[self.refndxmin:self.refndxmax].max())[0][-1] + self.refndxmin
            sd = self.max[n,0]+10*log10(exp(-1))
            self.skindx[n,0] = where(trace>sd)[0][-1]
            self.skindepth[n,0] = self.time[self.skindx[n,0]] # in ns
            self.refdepth[n,0] = self.time[self.refndx[n,0]] # in ns



def decay_map(range, **traits):
    from enthought.chaco.default_colormaps import ColorMapper
    """ Generator function for the decay colormap.
        The greenish color appears at approx 0.135 for a scale from 0.0 to 1.0
    """
    c = array([(0.7,0.1,0.0)]*255)
    c[34] = (0.0,0.5,0.2)
    c[0:34] = (0.0,0.1,0.7)
    colors = c
    return ColorMapper.from_palette_array(colors, range=range, **traits)


class DRange(Range):
    high = DelegatesTo('Dlow')
    Dlow = Instance(Range)
        
class TestTool(HasTraits):
    a = Range(-10, 10, 0)
    b = DRange(-10,10,0)
    b.Dlow = DelegatesTo("a")
    
    
    def _a_changed(self):
        print self.a
        
    def _b_changed(self):
        print self.b, self.a
    
if __name__ == "__main__":
    x = [1.2, 2.3, 3.5, 4.7, 6.3, 8.9, 10.6]
    y = [4.5, 8.7, 6.0, 2.9, 3.6, 4.5, 11.2]
    elev = [1000, 2000, 3500, 1700, 1300, 1000, 800]
    t = TrackPlot(x, y, elev)
    t.configure_traits()
#    spt = StartProcessingTool()
#    spt.configure_traits()